package cn.seecoder.ai.model.po;

import cn.seecoder.ai.dao.converter.TrainParamsConverter;
import cn.seecoder.ai.dao.converter.TrainResultsConverter;
import cn.seecoder.ai.enums.MachineLearningAlgorithmEnum;
import cn.seecoder.ai.enums.TrainParamEnum;
import cn.seecoder.ai.enums.TrainingStateEnum;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.hibernate.annotations.Comment;

import javax.persistence.*;
import java.util.Date;
import java.util.Map;

/**
 * 模型
 * @author fanyanpeng
 * @date 2023/4/9 22:14
 */

@Builder
@NoArgsConstructor
@AllArgsConstructor
@Data
@Entity
public class Model {

    @GeneratedValue(strategy = GenerationType.IDENTITY)
    @Id
    @Column(name = "id")
    private Integer id;

    @Basic
    @Column(name = "config_id")
    private Integer configId;

    @Basic
    @Column(name = "user_id")
    private Integer userId;

    @Basic
    @Column(name = "model_name")
    private String modelName;

    @Basic
    @Column(name = "machine_learning_algorithm")
    private MachineLearningAlgorithmEnum machineLearningAlgorithm;

    @Basic
    @Convert(converter = TrainParamsConverter.class)
    @Column(name = "train_params")
    private Map<TrainParamEnum,Object> trainParams;

    @Basic
    @Column(name = "training_state")
    @Enumerated(EnumType.STRING)
    @Comment("模型训练状态")
    private TrainingStateEnum trainingState;

    @Basic
    @Column(name = "train_result_map")
    @Convert(converter= TrainResultsConverter.class)
    @Comment("模型训练结果")
    private Map<String,Object> trainResults;

    @Basic
    @Column(name = "model_uri")
    @Comment("模型地址")
    private String modelUri;

    @Basic
    @Column(name = "features_indexer_uri")
    @Comment("训练模型时，需要将特征转换的转换器持久化，在模型使用时使用相同的转换器")
    private String featuresIndexerUri;

    @Basic
    @Column(name = "label_indexer_uri")
    @Comment("训练模型时，需要将标签转换的转换器持久化，在恢复标签时使用；标签为非数值类型时使用")
    private String labelIndexerUri;


    @Basic
    @Column(name = "training_start_time")
    @Comment("开始训练时间，点击训练时记录该时间")
    private Date trainingStartTime;

    @Basic
    @Column(name = "training_finish_time")
    @Comment("训练结束时记录该时间")
    private Date trainingFinishTime;

    @Basic
    @Column(name = "created_time")
    @Comment("模型创建时间")
    private Date createdTime;

    public Model setTrainParams(Map<TrainParamEnum, Object> trainParams) {
        this.trainParams = trainParams;
        return this;
    }

    public Map<TrainParamEnum, Object> getTrainParams() {
        return trainParams;
    }
}
